'''
Created on Oct 30, 2009

@author: mkiyer
'''

import veggie.coverage.covdb as covdb
from veggie.coverage.visualization.ucsc import get_bedgraph_track_header, gen_bedgraph_track_data
from veggie.io.bed import bed_reader
from veggie.sample.sampledb2 import get_sampledb
from veggie.sample.samplegroup import parse_samplegroups_xml
from veggie.gene.gene import parse_refseq_genes
from optparse import OptionParser
import collections
import h5py
import logging
import matplotlib.pyplot as plt
import numpy as np
import veggie.coverage.visualization.plots as covplot
import sys


def plot_utr_expr(chrom, start, end, strand, polya, sample_groups,
                  outfile=None,
                  **kwargs):
    '''
    kwargs:
    figsize - size of figure (tuple)
    title - title of plot (string)
    shade_start - instructs certain plots to shade a region
    shade_end - instructs certain plots to shade a region
    '''
    extended_utr = (polya, end)
    if strand == '-':
        extended_utr = (start, polya)
        
    # get coverage data for entire UTR
    utr_sample_coverages = covdb.calc_avg_cov_arrays(chrom, start, end, sample_groups)
    utr_summed_coverages = {}
    for s, covarray in utr_sample_coverages.iteritems():
        utr_summed_coverages[s] = covdb.calc_rpkm_array(covarray)

    # get coverage data for extended UTR
    extended_sample_coverages = covdb.calc_avg_cov_arrays(chrom, extended_utr[0], extended_utr[1], sample_groups)
    extended_summed_coverages = {}
    for s, covarray in extended_sample_coverages.iteritems():
        extended_summed_coverages[s] = covdb.calc_rpkm_array(covarray)

    # setup figure
    default_suptitle = '%s:%d-%d' % (chrom, start, end)
    suptitle = kwargs.get('title', default_suptitle)
    default_figsize = (20,10)
    #default_window_title = 'Boxplot'
    fig = plt.figure(figsize=kwargs.get('figsize', default_figsize))
    fig.suptitle(suptitle)
    #fig.canvas.set_window_title(kwargs.get('window_title', default_window_title))
    fig.subplots_adjust(left=0.05, right=0.975, top=0.95, bottom=0.125,
                        wspace=0.1, hspace=0.25)

    ax1 = fig.add_subplot(231)
    ax2 = fig.add_subplot(234)

    # Top plot shows average for entire UTR
    covplot.plot_sample_group_average_coverage(chrom, start, end, 
                                               sample_groups, 
                                               utr_sample_coverages,
                                               axes=ax1,
                                               **kwargs)
    # Bottom plot shows average for extended UTR
    covplot.plot_sample_group_average_coverage(chrom, 
                                               extended_utr[0], 
                                               extended_utr[1], 
                                               sample_groups, 
                                               extended_sample_coverages,
                                               axes=ax2,
                                               **kwargs)
    # Barplots
    ax3 = fig.add_subplot(232)
    ax4 = fig.add_subplot(235)
    covplot.plot_barplot(chrom, start, end,
                          sample_groups, utr_summed_coverages,
                          axes=ax3)
    covplot.plot_barplot(chrom, extended_utr[0], extended_utr[1],
                          sample_groups, extended_summed_coverages,
                          axes=ax4)
    # Boxplots
    ax5 = fig.add_subplot(233)
    ax6 = fig.add_subplot(236)
    covplot.plot_boxplot(chrom, start, end, 
                         sample_groups, utr_summed_coverages,
                         axes=ax5)
    covplot.plot_boxplot(chrom, start, end, 
                         sample_groups, extended_summed_coverages,
                         axes=ax6)

    # save to image
    if outfile is not None:
        plt.savefig(outfile)
    else:
        plt.savefig('%s_%d-%d.png' % (chrom, start, end))
    #plt.show()
    plt.close()
    return

if __name__ == '__main__':
    # configure logging
    logging.basicConfig(level=logging.DEBUG)

    # parse command line
    optionparser = OptionParser("usage: %prog [options] <regions.bed>")
    optionparser.add_option("--covdb", dest="covdb",
                            help="path to coverage database")
    optionparser.add_option("-s", "--samples", dest="samples",
                            help="sample group XML file")
    optionparser.add_option("-o", "--output", dest="outfile",
                            default=None,
                            help="output file [default: %default]")
    optionparser.add_option("-p", "--plots",
                            action="store_true", dest="plots", default=False,
                            help="create plots for each interval")
    optionparser.add_option("--bedgraph",
                            action="store_true", dest="bedgraph", default=False,
                            help="generate bedgraph track for UCSC browser")
    optionparser.add_option("--lanes", action="store_false", dest="pool_lanes",
                            default=True,
                            help="show individual lane results rather "
                            "than pooling")

    (options, args) = optionparser.parse_args()
    if options.outfile == None:
        outfhd = sys.stdout
    else:
        outfhd = open(options.outfile, 'w')
    if len(args) == 0:
        optionparser.error("no BED file specified")

    intervals = sorted(list(bed_reader(open(args[0]))))    
    # open coverage db
    h5file = h5py.File(options.covdb, 'r')
    # generate SampleGroupCoverage objects that pair samples with
    # coverage data
    sample_groups = [covdb.SampleGroupCoverage(sgroup, h5file) for sgroup in 
                     list(parse_samplegroups_xml(options.samples))]

    # process intervals
    for interval in intervals:
        title = '%s:%d-%d name=%s strand=%s' % \
            (interval.chrom, interval.start, interval.end, interval.name, interval.strand)
        outfile = '%s_%s_%d-%d.png' % (interval.name, interval.chrom, interval.start, interval.end) 
        logging.debug('plotting %s %s:%d-%d' % (interval.name, interval.chrom, interval.start, interval.end))
        plot_utr_expr(interval.chrom, interval.start, interval.end, interval.strand, interval.thickStart, 
                      sample_groups,
                      title=title,
                      outfile=outfile,
                      shade_start=interval.thickStart)
        
    h5file.close()